{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LangGraph Introduction\n",
    "\n",
    "This is a simple project for building ai agents with LangGraph and Amazon Bedrock Models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing necessary libraries\n",
    "\n",
    "from IPython.display import Image, display\n",
    "import gradio as gr\n",
    "from pydantic import BaseModel\n",
    "import random\n",
    "from typing import Annotated\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Using Amazon Bedrock Models\n",
    "from langchain_aws import ChatBedrockConverse \n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.graph import StateGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some useful constants to be randomly selected to form sentence\n",
    "\n",
    "nouns = [\"Cabbages\", \"Unicorns\", \"Toasters\", \"Penguins\", \"Bananas\", \"Zombies\", \"Rainbows\", \"Eels\", \"Pickles\", \"Muffins\"]\n",
    "adjectives = [\"outrageous\", \"smelly\", \"pedantic\", \"existential\", \"moody\", \"sparkly\", \"untrustworthy\", \"sarcastic\", \"squishy\", \"haunted\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading out environment variables\n",
    "# Using amazon bedrock, values like aws secret key, access key, bedrock api key and aws region name should exist\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A python function that capitalizes an input text. Using annotation for understanding\n",
    "# LangGraph uses reducer \n",
    "def shout(text: Annotated[str, \"Word to be capitalized\"]) -> str:\n",
    "    print(text.upper())\n",
    "\n",
    "shout(\"hello\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A word about \"Annotated\"\n",
    "\n",
    "You probably know this; type hinting is a feature in Python that lets you specify the type of something:\n",
    "\n",
    "`my_favorite_things: List`\n",
    "\n",
    "But you may not know this:\n",
    "\n",
    "You can also use something called \"Annotated\" to add extra information that somebody else might find useful:\n",
    "\n",
    "`my_favorite_things: Annotated[List, \"these are a few of mine\"]`\n",
    "\n",
    "LangGraph needs us to use this feature when we define our State object.\n",
    "\n",
    "It wants us to tell it what function it should call to update the State with a new value.\n",
    "\n",
    "This function is called a **reducer**.\n",
    "\n",
    "LangGraph provides a default reducer called `add_messages` which takes care of the most common case.\n",
    "\n",
    "And that hopefully explains why the State looks like this.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Define the State object\n",
    "\n",
    "You can use any python object; but it's most common to use a TypedDict or a Pydantic BaseModel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class State(BaseModel):\n",
    "    # Messages have the type \"list\". The `add_messages` function\n",
    "    # in the annotation defines how this state key should be updated\n",
    "    # (in this case, it appends messages to the list, rather than overwriting them)\n",
    "    messages: Annotated[list, add_messages]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Start the Graph Builder with this State class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_builder = StateGraph(State)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our graph can now handle two key tasks:\n",
    "\n",
    "- Each node can receive the current `State` as input and output an update to the state.\n",
    "- Updates to `messages` will be appended to the existing list rather than overwriting it, thanks to the prebuilt reducer function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Create a Node\n",
    "\n",
    "A node can be any python function.\n",
    "\n",
    "The reducer that we set before gets automatically called to combine this response with previous responses\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_sentence(old_state: State) -> State:\n",
    "\n",
    "    reply = f\"{random.choice(nouns)} are {random.choice(adjectives)}\"\n",
    "    messages = [{\"role\": \"assistant\", \"content\": reply}]\n",
    "\n",
    "    new_state = State(messages=messages)\n",
    "\n",
    "    return new_state\n",
    "\n",
    "graph_builder.add_node(\"first_node\", random_sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Create Edges\n",
    "\n",
    "- `entry` point to tell the graph where to start its work each time it is run\n",
    "- `exit` point to indicate where the graph should finish execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_builder.add_edge(START, \"first_node\")\n",
    "graph_builder.add_edge(\"first_node\", END)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Compile the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### That's it! Showtime!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(user_input: str, history):\n",
    "    message = {\"role\": \"user\", \"content\": user_input}\n",
    "    messages = [message]\n",
    "    state = State(messages=messages)\n",
    "    result = graph.invoke(state)\n",
    "    print(result)\n",
    "    return result[\"messages\"][-1].content\n",
    "\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangGraph is all about python functions - it doesn't need to involve LLMs!!\n",
    "\n",
    "Now we'll do the 5 steps again, but in 1 shot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Define the State object\n",
    "class State(BaseModel):\n",
    "    messages: Annotated[list, add_messages]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Start the Graph Builder with this State class\n",
    "graph_builder = StateGraph(State)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 3: Create a Node\n",
    "\n",
    "llm = ChatBedrockConverse(model_id=\"amazon.nova-micro-v1:0\")\n",
    "\n",
    "def chatbot_node(old_state: State) -> State:\n",
    "    response = llm.invoke(old_state.messages)\n",
    "    new_state = State(messages=[response])\n",
    "    return new_state\n",
    "\n",
    "graph_builder.add_node(\"chatbot\", chatbot_node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 4: Create Edges\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "graph_builder.add_edge(\"chatbot\", END)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 5: Compile the Graph\n",
    "graph = graph_builder.compile()\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the chatbot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(user_input: str, history):\n",
    "    initial_state = State(messages=[{\"role\": \"user\", \"content\": user_input}])\n",
    "    result = graph.invoke(initial_state)\n",
    "    print(result)\n",
    "    return result['messages'][-1].content\n",
    "\n",
    "\n",
    "gr.ChatInterface(chat, type=\"messages\").launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
